import argparse

from components.dataset import get_dataloader
from components.net import get_net
from components.loop import Loop
import torch
from torch.optim import SGD, Adam
from torch.nn import CrossEntropyLoss
from utils.random_seed import random_seed


def get_args():
    parser = argparse.ArgumentParser(description='Animals10 Project')
    parser.add_argument('--arch', default='resnet18', type=str)  # flag
    parser.add_argument('--epochs', default=300, type=int)

    parser.add_argument('--batch_size', default=256, type=int)
    parser.add_argument('--num_workers', default=8, type=int)
    parser.add_argument('--data_path', default='/home/heyujun/WorkSpace/Animals10/dataset/', type=str)

    parser.add_argument('--lr', default=0.1, type=float)  # flag

    parser.add_argument('--opti', default="SGD", type=str, help='[SGD, Adam]')

    parser.add_argument('--outputs', default='./outputs', type=str)
    parser.add_argument('--seed', default=0, type=int)

    return parser.parse_args()


def main():
    config = get_args()
    random_seed(config.seed)  # 随机种子

    train_loader, test_loader = get_dataloader(root=config.data_path,
                                               bs=config.batch_size,
                                               num_workers=config.num_workers)
    net = get_net(config.arch, num_classes=10)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    if config.opti == "SGD":
        opti = SGD(params=net.parameters(), lr=config.lr)
    elif config.opti == "Adam":
        opti = Adam(params=net.parameters(), lr=config.lr)

    loop = Loop(model=net.to(device),
                train_loader=train_loader,
                test_loader=test_loader,
                loss_fn=CrossEntropyLoss(),
                optimizer=opti,
                device=device)

    for epoch in range(1, config.epochs+1):
        loop.train(epoch)
        loop.test(epoch)
    loop.show()


if __name__ == "__main__":
    main()
